#include "rewine/impl/CPeImage.hpp"

#include "infra/Shelter.hpp"
#include "misc/arch.hpp"
#include "misc/helper.hpp"
#include "rewine/LoaderSingleton.hpp"
#include "rewine/LoggerSingleton.hpp"
#include "wininc/pe.h"

#include <sys/mman.h>
#include <wininc/pe.h>

using std::make_shared;
using std::shared_ptr;
using std::string;
using std::unique_ptr;

namespace rewine {

/**
 * relative virtual address (rva) to va
 */
PVOID
CPeImage::rva2va(LONG_PTR offset, const MapIndicator &map) const {
    // 0 is a special case
    if (!offset) return nullptr;
    if (offset >= map.size) {
        g_Logger->error(__FLF__, "rva2va: offset {} is out of {}", offset, map.size);
        return nullptr;
    }
    return (char *) map.start + offset;
}

PVOID
CPeImage::rva2va(LONG_PTR offset) const {
    return rva2va(offset, *m_spMap);
}

PVOID
CPeImage::rva2fp(PIMAGE_SECTION_HEADER secHdr, LONG_PTR offset, const MapIndicator &map) const {
    // 0 is a special case
    if (!offset) return nullptr;
    return (char *) map.start + (offset - secHdr->VirtualAddress + secHdr->PointerToRawData);
}

PVOID
CPeImage::rva2fp(PIMAGE_SECTION_HEADER secHdr, LONG_PTR offset) const {
    return rva2fp(secHdr, offset, *m_spSourceMap);
}

DWORD
CPeImage::directory_data_rva(IN USHORT index, OUT PULONG size) const {
    if (index >= opt_hdr->NumberOfRvaAndSizes) return 0;
    IMAGE_DATA_DIRECTORY *entry = opt_hdr->DataDirectory + index;
    if (!entry->VirtualAddress) return 0;
    if (size) *size = entry->Size;
    return entry->VirtualAddress;
}

void *
CPeImage::directory_data(IN USHORT index, OUT PULONG size) const {
    return rva2va(directory_data_rva(index, size));
}

int
CPeImage::parse_dos_header_and_goto_nt(MapIndicator &map, size_t offset) {
    auto szDos = sizeof(IMAGE_DOS_HEADER);
    PIMAGE_DOS_HEADER dos;
    if (offset == (size_t)-1) {
        dos = (PIMAGE_DOS_HEADER) map.move(0, szDos);
    } else {
        dos = (PIMAGE_DOS_HEADER) map.jump(offset, szDos);
    }
    if (!dos) {
        return -1;
    }
    if (dos->e_magic != IMAGE_DOS_SIGNATURE) {
        return -2;
    }
    map.jump(dos->e_lfanew, sizeof(IMAGE_NT_HEADERS));
    return 0;
}

int
CPeImage::use_nt_headers_and_next(MapIndicator &map, size_t offset) {
    auto szNt = sizeof(IMAGE_NT_HEADERS);
    PIMAGE_NT_HEADERS nt;
    if (offset == (size_t)-1) {
        nt = (PIMAGE_NT_HEADERS) map.move(0, szNt);
    } else {
        nt = (PIMAGE_NT_HEADERS) map.jump(offset, szNt);
    }
    if (!nt) {
        return -1;
    }
    map.next(0);
    if (nt->Signature != IMAGE_NT_SIGNATURE) {
        return -2;
    }
    nt_hdr = nt;
    file_hdr = &nt_hdr->FileHeader;
    opt_hdr = &nt_hdr->OptionalHeader;
    return 0;
}

int
CPeImage::check_headers_compatibility() {
    switch (file_hdr->Machine) {
        case IMAGE_FILE_MACHINE_I386:
            if (g_machine != machine_t::MACHINE_I386) {
                g_Logger->error(__FLF__, "i386 image but machine={}", to_string(g_machine));
                return -2;
            }
            break;
        case IMAGE_FILE_MACHINE_AMD64:
            if (g_machine != machine_t::MACHINE_AMD64) {
                g_Logger->error(__FLF__, "amd64 image but machine={}", to_string(g_machine));
                return -2;
            }
            break;
        default:
            g_Logger->error(__FLF__, "image unsupported, machine={}", file_hdr->Machine);
            return -2;
    }

    if (!file_hdr->SizeOfOptionalHeader) {
        return -3;
    }

    switch (opt_hdr->Magic) {
        case IMAGE_NT_OPTIONAL_HDR32_MAGIC:
            if (g_machine != machine_t::MACHINE_I386) return -4;
            break;
        case IMAGE_NT_OPTIONAL_HDR64_MAGIC:
            if (g_machine != machine_t::MACHINE_AMD64) return -4;
            break;
        default:
            return -4;
    }

    if (opt_hdr->DllCharacteristics & IMAGE_DLLCHARACTERISTICS_DYNAMIC_BASE) {
        // ASLR is good
    }

    bool is_flat = (opt_hdr->SectionAlignment & get_page_mask()) != 0;
    if (is_flat) {
        // not support yet
        return -6;
    }

    return 0;
}

int
CPeImage::use_section_headers_and_next(MapIndicator &map, size_t offset) {
    nbSection = file_hdr->NumberOfSections;
    auto szHeaders = nbSection * sizeof(IMAGE_SECTION_HEADER);
    if (offset == (size_t)-1) {
        sectionHeaders = (PIMAGE_SECTION_HEADER) map.move(0, szHeaders);
    } else {
        sectionHeaders = (PIMAGE_SECTION_HEADER) map.jump(offset, szHeaders);
    }
    if (!sectionHeaders) {
        return -1;
    }
    map.next(0);
    return 0;
}

PIMAGE_SECTION_HEADER
CPeImage::find_section_header_by_rva(DWORD vma) const {
    PIMAGE_SECTION_HEADER hdr = sectionHeaders;
    for (int i = 0; i < nbSection; ++i, ++hdr) {
        auto secStart = hdr->VirtualAddress;
        auto secEnd = secStart + hdr->Misc.VirtualSize;
        if (secStart <= vma && vma <= secEnd) {
            return hdr;
        }
    }
    return nullptr;
}

int
CPeImage::parse_dll_name() {
    auto exportd_rva = directory_data_rva(IMAGE_DIRECTORY_ENTRY_EXPORT, &szExportd);
    auto secHdr = find_section_header_by_rva(exportd_rva);
    exportd = (PIMAGE_EXPORT_DIRECTORY) rva2fp(secHdr, exportd_rva);
    if (!exportd) {
        return -1;
    }
    auto cname = (PCSTR) rva2fp(secHdr, exportd->Name);
    if (!cname) {
        return -2;
    }
    setExportName(string(cname));
    return 0;
}

int
mmap_image_base(void *preferred_start, size_t size, std::shared_ptr<MapIndicator> &map)
{
    map.reset();
    auto addr = ::mmap(preferred_start, size, PROT_NONE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    if (addr == (void *) -1) {
        return -1;
    }
    map = auto_MapIndicator(nullptr, 0, addr, size);
    return 0;
}

int
mmap_image_section(
    const std::shared_ptr<MapIndicator> &sourceMap, off_t offset, size_t copy_size,
    void *base, off_t rva, size_t mem_size, std::shared_ptr<MapIndicator> &mapSection)
{
    // for offset in PE is often unaligned with page, we just copy them to anonymous map instead
    mapSection.reset();

    // there may be extra data follows the loads
    if (copy_size > mem_size) copy_size = mem_size;

    auto start = (char *)base + rva;
    auto addr = ::mmap(start, mem_size, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED, -1, 0);
    if (addr == (void *) -1) {
        return -1;
    }

    memcpy(addr, (char *)sourceMap->start + offset, copy_size);
    if (copy_size < mem_size) {
        memset((char *)addr + copy_size, 0, mem_size - copy_size);
    }
    mapSection = auto_MapIndicator(sourceMap->pfd, offset, addr, mem_size);
    return 0;
}

int
CPeImage::map_base() {
    int ret;

    if (!m_bNeedRemap) {
        m_spMap = m_spSourceMap;
        return 0;
    }

    ret = mmap_image_base((void *)opt_hdr->ImageBase, opt_hdr->SizeOfImage, m_spMap);
    if (ret) {
        return -1;
    }
    return 0;
}

int
CPeImage::map_headers() {
    int ret;

    // copy DOS Header, NT Header, Section Headers
    auto szInFile = opt_hdr->SizeOfHeaders;
    ret = mmap_image_section(
        m_spSourceMap, 0, szInFile,
        m_spMap->front(), 0, szInFile, m_spHeadersMap);
    if (ret) {
        return -1;
    }

    ret = parse_dos_header_and_goto_nt(*m_spMap, 0);
    ret = use_nt_headers_and_next(*m_spMap);
    ret = use_section_headers_and_next(*m_spMap);
    return 0;
}

int
CPeImage::map_sections() {
    int ret;

    sectionMaps.clear();
    sectionMaps.reserve(nbSection);

    PIMAGE_SECTION_HEADER secHdr = sectionHeaders;
    for (int i = 0; i < nbSection; i++, secHdr++) {
        g_Logger->debug(
            __FLF__,
            "{}: {} - {} section {}",
            exportName().c_str(),
            (void *) ((char *) m_spMap->start + secHdr->VirtualAddress),
            (void *) ((char *) m_spMap->start + secHdr->VirtualAddress + secHdr->Misc.VirtualSize),
            secHdr->Name
        );

        std::shared_ptr<MapIndicator> spSecMap;
        // no check for PointerToRawData, becasuse we need to set protect flags
        // sometimes, SizeOfRawData is larger than VirtualSize, use VirtualSize in this case
        ret = mmap_image_section(
            m_spSourceMap, secHdr->PointerToRawData, secHdr->SizeOfRawData,
            m_spMap->front(), secHdr->VirtualAddress, secHdr->Misc.VirtualSize, spSecMap);
        if (ret) {
            return -1;
        }
        sectionMaps.push_back(spSecMap);
    }

    return 0;
}

static PIMAGE_BASE_RELOCATION
relocate_one(PVOID page, UINT count, USHORT *rel, INT_PTR delta) {
    INT_PTR loword = delta & 0xffff;
    INT_PTR hiword = (delta >> 16) & 0xffff;
    while (count--) {
        DWORD type = *rel >> 12;
        USHORT offset = *rel & 0xfff;
        PVOID addr = (char *) page + offset;
        //spdlog::info("type={} off={} addr={}", type, offset, addr);
        switch (type) {
            case IMAGE_REL_BASED_ABSOLUTE:
                break;
            case IMAGE_REL_BASED_HIGH:
                *(PWORD) addr += hiword;
                break;
            case IMAGE_REL_BASED_LOW:
                *(PWORD) addr += loword;
                break;
            case IMAGE_REL_BASED_HIGHLOW:
                *(PDWORD) addr += delta;
                break;
            case IMAGE_REL_BASED_HIGHADJ:
                rel++;
                count--;
                *(PDWORD) addr += (hiword << 16) + *rel;
                break;
            case IMAGE_REL_BASED_DIR64:
                *(PLONGLONG) addr += delta;
                break;
            case 6:
                // reserved
                break;
            default:
                g_Logger->error(__FLF__, "unsupported relocation. type={}", type);
        }
        rel++;
    }
    return (PIMAGE_BASE_RELOCATION) rel;  // return address of next block
}

int
CPeImage::relocate() {
    if (!(file_hdr->Characteristics & IMAGE_FILE_EXECUTABLE_IMAGE)) {
        g_Logger->warning(__FLF__, "[{}] non-executable, may be a linker error", fullname().c_str());
        return -1;
    }

    if (file_hdr->Characteristics & IMAGE_FILE_RELOCS_STRIPPED) {
        g_Logger->warning(__FLF__, "[{}] relocation info is stripped", fullname().c_str());
        return 0;
    }

    ULONG size;
    auto reld = (PIMAGE_BASE_RELOCATION) directory_data(IMAGE_DIRECTORY_ENTRY_BASERELOC, &size);
    if (!reld) return 0;

    INT_PTR delta = (INT_PTR) (m_spMap->start) - opt_hdr->ImageBase;
    //spdlog::info("{} map_start={} image_base={} delta={}", path.c_str(), image_map->start, opt_hdr->ImageBase, delta);

    size_t used = 0;
    while (used < size) {
        used += reld->SizeOfBlock;
        // reld is followed by some Type/Offset field entries
        UINT nb_block = (reld->SizeOfBlock - sizeof(IMAGE_BASE_RELOCATION)) / sizeof(WORD);
        // return the address of the next reld
        reld = relocate_one(rva2va(reld->VirtualAddress), nb_block, (PWORD) (reld + 1), delta);
    }

    return 0;
}

bool
CPeImage::is_forwarder_export(LONG_PTR rva) const {
    PVOID va = rva2va(rva);
    // if it is a forwarding export, va is a pointer to a NULL-term ASCII string in .edata section
    return (va >= (char *) exportd) && (va < (char *) exportd + szExportd);
}

int
CPeImage::load_exports() {
    exports.clear();

    exportd = (PIMAGE_EXPORT_DIRECTORY) directory_data(IMAGE_DIRECTORY_ENTRY_EXPORT, &szExportd);
    if (!exportd) return 0;

    nbExport = exportd->NumberOfFunctions;
    exports.reserve(nbExport);

    auto func_tbl = (PDWORD) rva2va(exportd->AddressOfFunctions);
    for (int i = 0; i < nbExport; i++) {
        auto pExport = CPeExport::create();
        exports.push_back(pExport);

        LONG_PTR rva = func_tbl[i];
        // we will process forwarder on bind_forwarder_exports() later
        bool bForward = is_forwarder_export(rva);
        PVOID va = rva2va(rva);
        pExport->setExport(i, bForward, va);

        g_Logger->trace(__FLF__, "@{} {}", i, pExport->getAddressDebugString());
    }

    // name is optional

    auto name_tbl = (PDWORD) rva2va(exportd->AddressOfNames);
    auto ordinal_tbl = (PWORD) rva2va(exportd->AddressOfNameOrdinals);
    if (name_tbl && ordinal_tbl) {
        auto ordinal_base = exportd->Base;
        for (int i = 0; i < exportd->NumberOfNames; ++i) {
            auto name = (PCSTR) rva2va(name_tbl[i]);
            DWORD ordinal = ordinal_base + ordinal_tbl[i];
            auto found = dynamic_cast<CPeExport *>(getExportByOrdinal(ordinal));
            if (found) {
                found->setName(i, name);
                g_Logger->trace(__FLF__, "@{} #{} hint={} {}", found->getIndex(), ordinal, i, name);
            } else {
                g_Logger->warning(__FLF__, "cannot find #{} hint={} {}", ordinal, i, name);
            }
        }
    }

    return 0;
}

static string
remove_tailing_spaces(PCSTR name) {
    if (!name) {
        return string();
    }
    size_t len = strlen(name);
    while (len && name[len - 1] == ' ') len--;
    return string(name, len);
}

int
CPeImage::load_imports() {
    nbImport = 0;
    imports.clear();

    ULONG size;
    auto impd_head = (PIMAGE_IMPORT_DESCRIPTOR) directory_data(IMAGE_DIRECTORY_ENTRY_IMPORT, &size);
    if (!impd_head) return 0;

    PIMAGE_IMPORT_DESCRIPTOR impd = impd_head;
    while (impd->u.Characteristics) {
        nbImport++;
        impd++;
    }
    if (!nbImport) return 0;

    imports.reserve(nbImport);

    impd = impd_head;
    for (int i = 0; i < nbImport; i++, impd++) {
        auto pDll = make_shared<CPeImportLib>();
        imports.push_back(pDll);
        load_import(i, impd, *pDll);
    }

    return 0;
}

int
CPeImage::load_import(int index, PIMAGE_IMPORT_DESCRIPTOR impd, CPeImportLib& dll) {
    int ret = 0;

    auto name = (PCSTR) rva2va(impd->Name);
    auto stripped_name = remove_tailing_spaces(name);
    dll.setName(index, stripped_name);

    g_Logger->debug(__FLF__, "{} import {}", exportName().c_str(), dll.getName().c_str());

    g_Loader->openImage(dll.getName(), false);

    // ignores BoundImportDirectory for it's just a cache on default
    // ilt for ImportLookupTable, iat for ImportAddressTable
    auto ilt = (PIMAGE_THUNK_DATA) rva2va(impd->u.OriginalFirstThunk); // ro
    auto iat = (PIMAGE_THUNK_DATA) rva2va(impd->FirstThunk); // rw

    PIMAGE_THUNK_DATA iltEntry, iatEntry;

    size_t nbSymbol = 0;
    iltEntry = ilt;
    while (iltEntry->u1.Ordinal) {
        nbSymbol++;
        iltEntry++;
    }
    dll.clearSymbols(nbSymbol);
    if (!nbSymbol) return 0;

    iltEntry = ilt;
    iatEntry = iat;
    for (int i = 0; i < nbSymbol; i++, iltEntry++, iatEntry++) {
        auto&& pImport = make_shared<CPeImport>();
        dll.addSymbol(pImport);
        ret = bind_import_symbol(dll, i, iltEntry, iatEntry, *pImport, true);
        if (ret < 0) {
            g_Logger->warning(
                __FLF__, "{} bind_import {} ({} {}) ret={}",
                exportName().c_str(),
                (void *) pImport->getSlotAddress(),
                dll.getName().c_str(), pImport->getName().c_str(),
                ret
            );
        } else {
            g_Logger->trace(
                __FLF__, "{} bind_import {} ({} {}) ret={}",
                exportName().c_str(),
                (void *) pImport->getSlotAddress(),
                dll.getName().c_str(), pImport->getName().c_str(),
                ret
            );
        }
    }

    return 0;
}

int
CPeImage::load_delayload_imports() {
    nbDelayload = 0;
    delayloads.clear();

    ULONG size;
    auto delay_head = (PIMAGE_DELAYLOAD_DESCRIPTOR) directory_data(IMAGE_DIRECTORY_ENTRY_DELAY_IMPORT, &size);
    if (!delay_head) return 0;

    PIMAGE_DELAYLOAD_DESCRIPTOR delayd = delay_head;
    while (delayd->DllNameRVA) {
        nbDelayload++;
        delayd++;
    }
    if (!nbDelayload) return 0;

    delayloads.reserve(nbDelayload);

    delayd = delay_head;
    for (int i = 0; i < nbDelayload; i++, delayd++) {
        auto pDll = make_shared<CPeImportLib>();
        delayloads.push_back(pDll);
        load_delayload_import(i, delayd, *pDll);
    }

    return 0;
}

int
CPeImage::load_delayload_import(int index, PIMAGE_DELAYLOAD_DESCRIPTOR delayd, CPeImportLib& dll) {
    int ret = 0;

    //TODO: handle delayd->Attributes

    auto name = (PCSTR) rva2va(delayd->DllNameRVA);
    auto stripped_name = remove_tailing_spaces(name);
    dll.setName(index, stripped_name);

    g_Logger->debug(__FLF__, "[{}] import_delayed {}", exportName().c_str(), dll.getName().c_str());

    g_Loader->openImage(dll.getName(), false);

    auto dnt = (PIMAGE_THUNK_DATA) rva2va(delayd->ImportNameTableRVA);
    auto dat = (PIMAGE_THUNK_DATA) rva2va(delayd->ImportAddressTableRVA);

    PIMAGE_THUNK_DATA dntEntry, datEntry;

    size_t nbSymbol = 0;
    dntEntry = dnt;
    while (dntEntry->u1.Ordinal) {
        nbSymbol++;
        dntEntry++;
    }
    dll.clearSymbols(nbSymbol);
    if (!nbSymbol) return 0;

    dntEntry = dnt;
    datEntry = dat;
    for (int i = 0; i < nbSymbol; i++, dntEntry++, datEntry++) {
        auto pImport = make_shared<CPeImport>();
        dll.addSymbol(pImport);
        ret = bind_import_symbol(dll, i, dntEntry, datEntry, *pImport, false);
        if (ret) {
            g_Logger->warning(
                __FLF__, "{} bind_delayload {} ({} {}) ret={}",
                exportName().c_str(),
                (void *) pImport->getSlotAddress(),
                dll.getName().c_str(), pImport->getName().c_str(),
                ret
            );
        } else {
            g_Logger->trace(
                __FLF__, "{} bind_delayload {} ({} {}) ret={}",
                exportName().c_str(),
                (void *) pImport->getSlotAddress(),
                dll.getName().c_str(), pImport->getName().c_str(),
                ret
            );
        }
    }

    return 0;
}

int
CPeImage::bind_import_symbol(CPeImportLib& dll, int index, PIMAGE_THUNK_DATA pName, PIMAGE_THUNK_DATA pAddress, CPeImport& imp, bool bBindNow) {
    imp.setImport(index, pName, pAddress);

    if (IMAGE_SNAP_BY_ORDINAL(pName->u1.Ordinal)) {
        DWORD ordinal = IMAGE_ORDINAL(pName->u1.Ordinal);
        imp.setOrdinal(ordinal);
    } else {
        auto info = (PIMAGE_IMPORT_BY_NAME) rva2va(pName->u1.AddressOfData);
        imp.setHintAndName(info->Hint, string(info->Name));
    }

    // notice: getName will try using ordinal first
    g_Logger->trace(__FLF__, "resolve import: {} {}", dll.getName(), imp.getName());

    if (g_Loader->getSymbolResolver().resolveAndBind(this, dll.getName(), imp.getName(), imp.getHint(), imp.getValueSynchronizerRef(), bBindNow)) {
        return 0;
    }

    //TODO set stub
    return -1;
}

int
CPeImage::bind_forwarder_exports() {
    int ret;
    for (auto& pExport : exports) {
        if (!pExport->isForwarder()) continue;
        ret = bind_forwarder_export(*pExport);
        if (ret) {
            g_Logger->warning(
                __FLF__, "{} bind_forwarder_export {} ret={}",
                exportName().c_str(),
                pExport->getName().c_str(),
                ret
            );
        }
    }
    return 0;
}

int
CPeImage::bind_forwarder_export(CPeExport& exp) {
    auto forwarder = exp.getForwarder();

    g_Logger->trace(__FLF__, "resolve forwardTo: \"{}\"", forwarder);

    auto pos = forwarder.rfind('.');
    if (pos == string::npos) {
        g_Logger->warning(__FLF__, "broken forward string: \"{}\"", forwarder);
        return -1;
    }

    string dllName = forwarder.substr(0, pos) + ".dll";
    string symbolName = forwarder.substr(pos + 1);

    if (g_Loader->getSymbolResolver().resolveAndBind(this, dllName, symbolName, -1, exp.getValueSynchronizerRef(), false)) {
        return 0;
    }

    //TODO set stub?
    return -2;
}

int
CPeImage::set_memory_acl(void) {
    mprotect(m_spMap->front(), m_spMap->size, PROT_READ);

    PIMAGE_SECTION_HEADER secHdr = sectionHeaders;

    for (int i = 0; i < nbSection; i++, secHdr++) {
        DWORD charact = secHdr->Characteristics;
        int flags = 0;
        if (charact & IMAGE_SCN_MEM_READ) flags |= PROT_READ;
        if (charact & IMAGE_SCN_MEM_WRITE) flags |= PROT_WRITE;
        if (charact & IMAGE_SCN_MEM_EXECUTE) flags |= PROT_EXEC;

        auto& apSecMap = sectionMaps[i];
        mprotect(apSecMap->front(), apSecMap->size, flags);
    }

    ULONG dir_size;
    PVOID iat = directory_data(IMAGE_DIRECTORY_ENTRY_IAT, NULL);
    //TODO allow write

    return 0;
}

int
CPeImage::initFromMap(const std::shared_ptr<MapIndicator> &map, bool bRemap) {
    m_spSourceMap = map;
    m_spMap = nullptr;
    m_bNeedRemap = bRemap;
    return 0;
}

int
CPeImage::initFromFd(int fd) {
    int ret;
    if (fd < 0) {
        return -1;
    }
    auto pfd = auto_fd(fd);
    if (!pfd) {
        return -2;
    }
    shared_ptr<MapIndicator> apMap;
    ret = mmap_readonly_file(pfd, apMap);
    if (ret < 0) {
        return -3;
    }
    return initFromMap(apMap, true);
}

int CPeImage::parse() {
    int ret;

    ret = parse_dos_header_and_goto_nt(*m_spSourceMap, 0);
    if (ret < 0) {
        g_Logger->error(__FLF__, "{} parse_dos_header ret={}", fullname().c_str(), ret);
        return -1;
    }

    ret = use_nt_headers_and_next(*m_spSourceMap);
    if (ret < 0) {
        g_Logger->error(__FLF__, "{} use_nt_headers ret={}", fullname().c_str(), ret);
        return -2;
    }

    ret = check_headers_compatibility();
    if (ret < 0) {
        g_Logger->error(__FLF__, "{} check_headers_compatibility ret={}", fullname().c_str(), ret);
        return -3;
    }

    ret = use_section_headers_and_next(*m_spSourceMap);
    if (ret < 0) {
        g_Logger->error(__FLF__, "{} use_section_headers ret={}", fullname().c_str(), ret);
        return -4;
    }

    ret = parse_dll_name();
    if (ret < 0) {
        g_Logger->error(__FLF__, "{} parse_dll_name ret={}", fullname().c_str(), ret);
        return -5;
    }

    return 0;
}

int CPeImage::load() {
    int ret;

    setIsLoading();
    Shelter shlLoadStatus;
    shlLoadStatus.add([this]() {
        this->setIsLoaded(false);
    });

    ret = map_base();
    if (ret < 0) {
        g_Logger->error(__FLF__, "{} map_base ret={}", fullname().c_str(), ret);
        return -1;
    }
    ret = map_headers();
    if (ret < 0) {
        g_Logger->error(__FLF__, "{} map_headers ret={}", fullname().c_str(), ret);
        return -1;
    }
    ret = map_sections();
    if (ret < 0) {
        g_Logger->error(__FLF__, "{} map_sections ret={}", exportName().c_str(), ret);
        return -1;
    }
    // finished
    m_spSourceMap.reset();

    ret = load_exports();
    if (ret < 0) {
        g_Logger->error(__FLF__, "{} load_exports ret={}", exportName().c_str(), ret);
        return -2;
    }

    ret = relocate();
    if (ret < 0) {
        g_Logger->error(__FLF__, "{} relocate ret={}", exportName().c_str(), ret);
        return -3;
    }

    ret = load_imports();
    if (ret < 0) {
        g_Logger->error(__FLF__, "{} load_imports ret={}", exportName().c_str(), ret);
        return -4;
    }

    ret = load_delayload_imports();
    if (ret < 0) {
        g_Logger->error(__FLF__, "{} load_delayload_imports ret={}", exportName().c_str(), ret);
        return -5;
    }

    ret = bind_forwarder_exports();
    if (ret < 0) {
        g_Logger->error(__FLF__, "{} bind_forwarder_exports ret={}", exportName().c_str(), ret);
        return -6;
    }

    ret = set_memory_acl();
    if (ret < 0) {
        g_Logger->error(__FLF__, "{} set_memory_acl ret={}", exportName().c_str(), ret);
        return -7;
    }

    shlLoadStatus.toggle(false);
    return 0;
}

int CPeImage::unload() {
    imports.clear();
    nbImport = 0;

    delayloads.clear();
    nbDelayload = 0;

    exports.clear();
    nbExport = 0;

    szExportd = 0;
    exportd = nullptr;

    sectionMaps.clear();
    nbSection = 0;
    sectionHeaders = nullptr;

    nt_hdr = nullptr;
    file_hdr = nullptr;
    opt_hdr = nullptr;

    if (m_bNeedRemap) {
        m_spMap.reset();
    }
    m_spSourceMap.reset();

    setIsUnloaded();
    return 0;
}

int
CPeImage::getExportCount() const {
    return nbExport;
}

IExport *
CPeImage::getExportByIndex(const unsigned int index) const {
    if (index < 0 || index >= nbExport) {
        return nullptr;
    }
    return exports[index].get();
}

IExport *
CPeImage::getExportByOrdinal(DWORD ordinal) const {
    int index = ordinal - exportd->Base;
    return getExportByIndex(index);
}

struct bsearch_export_nametbl_context_t {
    const CPeImage *pImage;
    const char *name;
};

static int
bsearch_export_nametbl(const void *ctx0, const void *ele) {
    struct bsearch_export_nametbl_context_t *ctx = (struct bsearch_export_nametbl_context_t *) ctx0;
    auto name = (const char *) ctx->pImage->rva2va(*reinterpret_cast<const DWORD *>(ele));
    return strcmp(ctx->name, name);
}

IExport *
CPeImage::getExportByName(const string& name, uint16_t hint) const {
    if (!name.empty()) {
        if (name[0] == '#') {
            DWORD ordinal = (DWORD) std::stoi(name.substr(1));
            return getExportByOrdinal(ordinal);
        }
    }

    //TODO: improve: try name_tbl[hint] first if hint  [0, nb_name), bsearch if fails
    auto iter = std::find_if(exports.cbegin(), exports.cend(), [&name](const shared_ptr<CExport>& exp) -> bool {
        return name == exp->getName();
    });
    if (iter != exports.cend()) {
        return (*iter).get();
    }
    return nullptr;
    /*
    DWORD name_index;
    auto name_tbl = (PDWORD) rva2va(exportd->AddressOfNames);
    if (hint < exportd->NumberOfNames && strcmp(name.c_str(), (const char *)rva2va(name_tbl[hint])) == 0) {
        name_index = hint;
    } else {
        struct bsearch_export_nametbl_context_t ctx;
        ctx.pImage = this;
        ctx.name = name.c_str();
        auto pFound = (PDWORD) bsearch(&ctx, name_tbl, exportd->NumberOfNames, sizeof(DWORD), bsearch_export_nametbl);
        if (!pFound) {
            return nullptr;
        }
        name_index = pFound - name_tbl;
    }
    auto ordinal_tbl = (PWORD) rva2va(exportd->AddressOfNameOrdinals);
    auto ordinal = ordinal_tbl[name_index];
    return getExportByIndex(ordinal);
     */
}

BOOL
CPeImage::onEvent(DWORD dwReason) const {
    auto dllmain = (dllmain_t) getEntryPoint();
    return dllmain(getImageBase(), dwReason, nullptr);
}

} // namespace rewine
